"""
抖音视频搜索脚本
用于根据关键词搜索抖音视频并提取视频信息
"""

import asyncio
import sys
import os
import json
import re

# Add the project root directory to the Python path
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from playwright.async_api import async_playwright


class DouyinVideoSearcher:
    """
    抖音视频搜索器，用于搜索和提取抖音视频信息
    """
    
    def __init__(self):
        """
        初始化抖音视频搜索器
        """
        self.playwright = None
        self.context = None
        self.page = None

    async def start_browser(self, headless=False):
        """
        启动浏览器
        """
        self.playwright = await async_playwright().start()
        self.context = await self.playwright.chromium.launch_persistent_context(
            user_data_dir="./user_data/douyin",
            channel="chrome",
            headless=headless,
            args=[
                "--no-sandbox",
                "--disable-dev-shm-usage",
                "--disable-blink-features=AutomationControlled",
                "--disable-extensions"
            ],
            viewport={"width": 1920, "height": 1080}
        )
        self.page = await self.context.new_page()
        return self.context, self.page

    async def search_videos(self, keyword, max_count=50):
        """
        搜索抖音视频并提取信息
        
        Args:
            keyword (str): 搜索关键词
            max_count (int): 最大获取视频数量
            
        Returns:
            list: 视频信息列表
        """
        # 导航到抖音搜索页面
        search_url = f"https://www.douyin.com/search/{keyword}?type=video"
        await self.page.goto(search_url)
        await self.page.wait_for_timeout(3000)  # 等待页面加载
        
        # 滚动页面以加载更多视频
        videos = []
        scroll_count = 0
        max_scroll_count = max_count // 10 + 3  # 根据需要的视频数量估算滚动次数
        
        while len(videos) < max_count and scroll_count < max_scroll_count:
            # 执行滚动
            await self.page.evaluate("window.scrollTo(0, document.body.scrollHeight)")
            await self.page.wait_for_timeout(2000)  # 等待内容加载
            scroll_count += 1
            
            # 提取当前页面的视频信息
            current_videos = await self.page.evaluate(r'''() => {
                const videos = [];
                // 使用更精确的选择器查找视频元素
                let videoItems = Array.from(document.querySelectorAll('div[data-e2e="search-result-container"] div a[href*="/video/"]'));
                
                // 如果上面的选择器没有找到元素，尝试备用选择器
                if (videoItems.length === 0) {
                    const backupItems = document.querySelectorAll('a[href*="/video/"]');
                    videoItems = Array.from(backupItems).filter(item => item.href.includes('/video/'));
                }
                
                videoItems.forEach((item, index) => {
                    if (index >= 100) return; // 只处理前100个视频
                    
                    const url = item.href;
                    // 从URL中提取aweme_id
                    const urlParts = url.split('/');
                    const awemeId = urlParts[urlParts.length - 1];
                    
                    const textContent = item.textContent || '';
                    
                    // 提取用户名（以@开头的文本）
                    let username = '未知用户';
                    const usernameMatch = textContent.match(/@[^\\s]+/);
                    if (usernameMatch) {
                        username = usernameMatch[0];
                    }
                    
                    // 提取观看数（包含数字和单位的文本）
                    let views = '观看数未知';
                    const viewsMatch = textContent.match(/\\d+[.\\d]*\\s*[万]*/);
                    if (viewsMatch) {
                        views = viewsMatch[0];
                    }
                    
                    // 提取标题（查找包含中文和标点的文本）
                    let title = '';
                    // 查找更具体的标题元素
                    const titleElement = item.querySelector('img[alt]') || item.querySelector('img');
                    if (titleElement && titleElement.alt) {
                        title = titleElement.alt.trim();
                    } else if (textContent) {
                        // 从文本内容中提取标题
                        const titlePattern = /[\\u4e00-\\u9fa5]+[^\\n]*?[!?！？]*[\\u4e00-\\u9fa5]/;
                        const match = textContent.match(titlePattern);
                        if (match) {
                            title = match[0].trim();
                        }
                    }
                    
                    // 如果仍然没有标题，则使用部分文本内容
                    if (!title && textContent) {
                        title = textContent.substring(0, 50).trim();
                    }
                    
                    videos.push({
                        aweme_id: awemeId,
                        url: url,
                        title: title,
                        username: username,
                        views: views,
                        full_title: textContent
                    });
                });
                
                return videos;
            }''')
            
            # 合并视频信息并去重
            for video in current_videos:
                # 检查是否已存在该视频
                if not any(v['aweme_id'] == video['aweme_id'] for v in videos):
                    videos.append(video)
            
            if len(videos) >= max_count:
                break
        
        # 限制返回的视频数量
        return videos[:max_count]
    
    async def close(self):
        """
        关闭浏览器
        """
        if self.context:
            await self.context.close()
        if self.playwright:
            await self.playwright.stop()


async def main():
    """
    主函数，处理命令行参数并执行搜索
    """
    if len(sys.argv) < 2:
        print("使用方法: python douyin_video_search.py <关键词> [数量]")
        print("示例: python douyin_video_search.py 足球高清 50")
        return
    
    keyword = sys.argv[1]
    max_count = 50  # 默认获取50个视频
    
    if len(sys.argv) > 2:
        try:
            max_count = int(sys.argv[2])
        except ValueError:
            print("数量参数必须是整数，使用默认值50")
    
    print(f"正在搜索关键词: {keyword}, 需要获取 {max_count} 个视频")
    
    # 创建搜索器实例
    searcher = DouyinVideoSearcher()
    
    try:
        # 启动浏览器
        await searcher.start_browser(headless=False)
        
        # 搜索视频
        videos = await searcher.search_videos(keyword, max_count)
        
        # 输出JSON格式结果
        print(json.dumps(videos, ensure_ascii=False, indent=2))
            
    except Exception as e:
        print(f"搜索过程中发生错误: {e}")
    finally:
        # 关闭浏览器
        await searcher.close()


if __name__ == "__main__":
    asyncio.run(main())